import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
import json
from models import Encoder, DecoderWithAttention
import cv2


class ImageCaption:
    def __init__(self, encoder_path, decoder_path, word_map_path):
        # parameter
        # Model parameters
        emb_dim = 512  # dimension of word embeddings
        attention_dim = 512  # dimension of attention linear layers
        decoder_dim = 512  # dimension of decoder RNN
        dropout = 0.5
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # self.device = torch.device("cpu")

        # Load models
        self.encoder = Encoder()
        self.decoder = DecoderWithAttention(
            attention_dim=attention_dim,
            embed_dim=emb_dim,
            decoder_dim=decoder_dim,
            vocab_size=len(self.load_word_map(word_map_path)),
            dropout=dropout,
        )

        # Load weights
        self.encoder.load_state_dict(
            torch.load(
                encoder_path, map_location=torch.device("cpu"), weights_only=True
            )
        )
        self.decoder.load_state_dict(
            torch.load(
                decoder_path, map_location=torch.device("cpu"), weights_only=True
            )
        )

        # Set to evaluation mode
        self.encoder.eval()
        self.decoder.eval()

        # Load word map
        self.word_map = self.load_word_map(word_map_path)
        self.rev_word_map = {v: k for k, v in self.word_map.items()}  # ix2word

    def load_word_map(self, path):
        with open(path, "r") as j:
            word_map = json.load(j)
        return word_map

    def caption_image_beam_search(self, img, beam_size=5):
        # Preprocess image
        k = beam_size
        
        img = cv2.resize(img, (256, 256))
        img = img.transpose(2, 0, 1)
        img = torch.FloatTensor(img).to(self.device) / 255.0
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        )
        transform = transforms.Compose([normalize])
        image = transform(img).unsqueeze(0)  # (1, 3, 256, 256)

        # Encode
        encoder_out = self.encoder(image)
        enc_image_size = encoder_out.size(1)
        encoder_dim = encoder_out.size(3)

        # Flatten encoding
        encoder_out = encoder_out.view(1, -1, encoder_dim)
        num_pixels = encoder_out.size(1)

        # We'll treat the problem as having a batch size of k
        encoder_out = encoder_out.expand(beam_size, num_pixels, encoder_dim)

        # Tensor to store top k previous words at each step; now they're just <start>
        k_prev_words = torch.LongTensor([[self.word_map["<start>"]]] * beam_size).to(
            self.device
        )

        # Tensor to store top k sequences; now they're just <start>
        seqs = k_prev_words  # (k, 1)

        # Tensor to store top k sequences' scores; now they're just 0
        top_k_scores = torch.zeros(beam_size, 1).to(self.device)  # (k, 1)

        # Tensor to store top k sequences' alphas; now they're just 1s
        seqs_alpha = torch.ones(beam_size, 1, enc_image_size, enc_image_size).to(
            self.device
        )

        # Lists to store completed sequences, their alphas and scores
        complete_seqs = list()
        complete_seqs_alpha = list()
        complete_seqs_scores = list()

        # Start decoding
        step = 1
        h, c = self.decoder.init_hidden_state(encoder_out)

        # s is a number less than or equal to k, because sequences are removed from this process once they hit <end>
        while True:
            embeddings = self.decoder.embedding(k_prev_words).squeeze(1)
            awe, alpha = self.decoder.attention(encoder_out, h)
            alpha = alpha.view(
                -1, enc_image_size, enc_image_size
            )  # (s, enc_image_size, enc_image_size)
            gate = self.decoder.sigmoid(self.decoder.f_beta(h))
            awe = gate * awe
            h, c = self.decoder.decode_step(torch.cat([embeddings, awe], dim=1), (h, c))
            scores = self.decoder.fc(h)  # (s, vocab_size)
            scores = F.log_softmax(scores, dim=1)

            # Add
            scores = top_k_scores.expand_as(scores) + scores  # (s, vocab_size)

            # For the first step, all k points will have the same scores (since same k previous words, h, c)
            if step == 1 or beam_size >= scores.size(0):
                top_k_scores, top_k_words = scores[0].topk(beam_size, 0, True, True)
            else:
                # Unroll and find top scores, and their unrolled indices
                top_k_scores, top_k_words = scores.view(-1).topk(
                    beam_size, 0, True, True
                )

            # Convert unrolled indices to actual indices of scores
            prev_word_inds = top_k_words / len(self.word_map)
            next_word_inds = top_k_words % len(self.word_map)

            # Add new words to sequences, alphas
            prev_word_inds = prev_word_inds.long()
            seqs = torch.cat([seqs[prev_word_inds], next_word_inds.unsqueeze(1)], dim=1)
            seqs_alpha = torch.cat(
                [seqs_alpha[prev_word_inds], alpha[prev_word_inds].unsqueeze(1)], dim=1
            )

            # Which sequences are incomplete (didn't reach <end>)?
            incomplete_inds = [
                ind
                for ind, next_word in enumerate(next_word_inds)
                if next_word != self.word_map["<end>"]
            ]
            complete_inds = list(set(range(len(next_word_inds))) - set(incomplete_inds))

            # Set aside complete sequences
            if len(complete_inds) > 0:
                complete_seqs.extend(seqs[complete_inds].tolist())
                complete_seqs_alpha.extend(seqs_alpha[complete_inds].tolist())
                complete_seqs_scores.extend(top_k_scores[complete_inds])
            k -= len(complete_inds)  # reduce beam length accordingly

            # Proceed with incomplete sequences
            if k == 0:
                break
            seqs = seqs[incomplete_inds]
            seqs_alpha = seqs_alpha[incomplete_inds]
            h = h[prev_word_inds[incomplete_inds]]
            c = c[prev_word_inds[incomplete_inds]]
            encoder_out = encoder_out[prev_word_inds[incomplete_inds]]
            top_k_scores = top_k_scores[incomplete_inds].unsqueeze(1)
            k_prev_words = next_word_inds[incomplete_inds].unsqueeze(1)

            # Break if things have been going on too long
            if step > 50:
                break
            step += 1

        i = complete_seqs_scores.index(max(complete_seqs_scores))
        seq = complete_seqs[i]
        alphas = complete_seqs_alpha[i]

        words = [self.rev_word_map[ind] for ind in seq]
        sentence = " ".join(words[1:-1])
        return sentence, alphas

    def caption(self, img):
        sentence, _ = self.caption_image_beam_search(img)
        return sentence


if __name__ == "__main__":
    modeldir = "models/"
    encoder_path = modeldir + "11encoder_weights.pth"
    decoder_path = modeldir + "11decoder_weights.pth"
    word_map_path = modeldir + "WORDMAP_coco_5_cap_per_img_5_min_word_freq.json"

    print("loading model.")
    image_caption = ImageCaption(encoder_path, decoder_path, word_map_path)

    print('test begin')
    img = cv2.imread("testImage/dog.jpeg")
    caption = image_caption.caption(img)
    print("Generated Caption:", caption)
